// Copyright 2019 StreamSets Inc.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package iothub

import (
	"bytes"
	"context"
	"github.com/amenzhinsky/iothub/iotdevice"
	iotmqtt "github.com/amenzhinsky/iothub/iotdevice/transport/mqtt"
	"github.com/sirupsen/logrus"
	"github.com/streamsets/datacollector-edge/api"
	"github.com/streamsets/datacollector-edge/api/dataformats"
	"github.com/streamsets/datacollector-edge/api/validation"
	"github.com/streamsets/datacollector-edge/container/common"
	"github.com/streamsets/datacollector-edge/stages/lib/datagenerator"
	"github.com/streamsets/datacollector-edge/stages/stagelibrary"
)

const (
	Library   = "streamsets-datacollector-azure-lib"
	StageName = "com_streamsets_pipeline_stage_destination_iothub_IotHubProducerDTarget"
)

type Destination struct {
	*common.BaseStage
	ProducerConf ProducerConfigBean `ConfigDefBean:"producerConf"`
	iotHubClient *iotdevice.Client
}

type ProducerConfigBean struct {
	IotHubName                string                                  `ConfigDef:"type=STRING,required=true"`
	DeviceId                  string                                  `ConfigDef:"type=STRING,required=true"`
	SasKey                    string                                  `ConfigDef:"type=STRING,required=true"`
	MaxRequestCompletionSecs  float64                                 `ConfigDef:"type=NUMBER,required=true"`
	DataFormat                string                                  `ConfigDef:"type=STRING,required=true"`
	DataGeneratorFormatConfig datagenerator.DataGeneratorFormatConfig `ConfigDefBean:"dataGeneratorFormatConfig"`
}

func init() {
	stagelibrary.SetCreator(Library, StageName, func() api.Stage {
		return &Destination{BaseStage: &common.BaseStage{}}
	})
}

func (dest *Destination) Init(stageContext api.StageContext) []validation.Issue {
	issues := dest.BaseStage.Init(stageContext)
	var err error
	dest.iotHubClient, err = iotdevice.New(
		iotdevice.WithTransport(iotmqtt.New()),
		iotdevice.WithConnectionString(dest.ProducerConf.SasKey),
	)
	if err != nil {
		issues = append(issues, stageContext.CreateConfigIssue(err.Error()))
		return issues
	}
	// connect to the iothub
	if err := dest.iotHubClient.Connect(context.Background()); err != nil {
		issues = append(issues, stageContext.CreateConfigIssue(err.Error()))
		return issues
	}
	issues = dest.ProducerConf.DataGeneratorFormatConfig.Init(dest.ProducerConf.DataFormat, stageContext, issues)
	return issues
}

func (dest *Destination) Write(batch api.Batch) error {
	records := batch.GetRecords()
	recordWriterFactory := dest.ProducerConf.DataGeneratorFormatConfig.RecordWriterFactory
	for _, record := range records {
		recordBuffer := bytes.NewBuffer([]byte{})
		recordWriter, err := recordWriterFactory.CreateWriter(dest.GetStageContext(), recordBuffer)
		if err != nil {
			logrus.WithError(err).Error("Error creating writer")
			dest.GetStageContext().ToError(err, record)
			break
		}
		err = recordWriter.WriteRecord(record)
		if err != nil {
			logrus.WithError(err).Error("Error writing record")
			dest.GetStageContext().ToError(err, record)
			break
		}
		flushAndCloseWriter(recordWriter)

		// send a device-to-cloud message
		if err := dest.iotHubClient.SendEvent(context.Background(), recordBuffer.Bytes()); err != nil {
			logrus.WithError(err).Error("Error sending event to Azure IoT Hub")
			dest.GetStageContext().ToError(err, record)
		}
	}

	return nil
}

func flushAndCloseWriter(recordWriter dataformats.RecordWriter) {
	err := recordWriter.Flush()
	if err != nil {
		logrus.WithError(err).Error("Error flushing record writer")
	}

	err = recordWriter.Close()
	if err != nil {
		logrus.WithError(err).Error("Error closing record writer")
	}
}
